package com.sinosoft.common.encrypt.filter;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.sinosoft.common.encrypt.properties.ApiDecryptProperties;
import com.sinosoft.common.encrypt.utils.CryptoUtils;
import jakarta.servlet.ServletOutputStream;
import jakarta.servlet.WriteListener;
import jakarta.servlet.http.HttpServletResponse;
import jakarta.servlet.http.HttpServletResponseWrapper;
import java.nio.charset.StandardCharsets;

import java.io.*;

/**
 * 加密响应体包装器
 */
public class EncryptResponseBodyWrapper extends HttpServletResponseWrapper {
    
    private final ApiDecryptProperties properties;
    private final ByteArrayOutputStream buffer = new ByteArrayOutputStream();
    private final ServletOutputStream outputStream;
    private final PrintWriter writer;
    
    public EncryptResponseBodyWrapper(HttpServletResponse response, ApiDecryptProperties properties) {
        super(response);
        this.properties = properties;
        this.outputStream = new ServletOutputStream() {
            @Override
            public void write(int b) throws IOException {
                buffer.write(b);
            }
            
            @Override
            public boolean isReady() {
                return false;
            }
            
            @Override
            public void setWriteListener(WriteListener writeListener) {
            }
        };
        this.writer = new PrintWriter(new OutputStreamWriter(buffer, StandardCharsets.UTF_8));
    }
    
    @Override
    public ServletOutputStream getOutputStream() throws IOException {
        return outputStream;
    }
    
    @Override
    public PrintWriter getWriter() throws IOException {
        return writer;
    }
    
    public String getEncryptContent(HttpServletResponse response, String publicKey, String headerFlag) throws IOException {
        // 获取原始响应内容
        String originalContent = buffer.toString("UTF-8");
        
        // 生成SM4密钥
        String sm4Key = properties.getSk();
        
        // 使用SM4密钥加密响应内容
        String encryptedContent = CryptoUtils.encryptBody(originalContent, sm4Key);
        
        // 使用SM2公钥加密SM4密钥
        String encryptedKey = CryptoUtils.encryptSm4Key(sm4Key, properties.getSm2PublicKey());
        
        // 设置加密密钥头
        response.setHeader(headerFlag, encryptedKey);
        
        return encryptedContent;
    }
} 